{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "fe44d03a",
   "metadata": {},
   "source": [
    "# Custom DataLoader for Imbalanced dataset"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5731376a",
   "metadata": {},
   "source": [
    "* In this notebook we will use the higly imbalanced Protein Homology Dataset from [KDD cup 2004](https://www.kdd.org/kdd-cup/view/kdd-cup-2004/Data)\n",
    "\n",
    "```\n",
    "* The first element of each line is a BLOCK ID that denotes to which native sequence this example belongs. There is a unique BLOCK ID for each native sequence. BLOCK IDs are integers running from 1 to 303 (one for each native sequence, i.e. for each query). BLOCK IDs were assigned before the blocks were split into the train and test sets, so they do not run consecutively in either file.\n",
    "* The second element of each line is an EXAMPLE ID that uniquely describes the example. You will need this EXAMPLE ID and the BLOCK ID when you submit results.\n",
    "* The third element is the class of the example. Proteins that are homologous to the native sequence are denoted by 1, non-homologous proteins (i.e. decoys) by 0. Test examples have a \"?\" in this position.\n",
    "* All following elements are feature values. There are 74 feature values in each line. The features describe the match (e.g. the score of a sequence alignment) between the native protein sequence and the sequence that is tested for homology.\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8969d2a7",
   "metadata": {},
   "source": [
    "## Initial imports"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "5d189eca",
   "metadata": {},
   "outputs": [],
   "source": [
    "from pytorch_widedeep import Trainer\n",
    "from pytorch_widedeep.preprocessing import TabPreprocessor\n",
    "from pytorch_widedeep.models import TabMlp, WideDeep\n",
    "from pytorch_widedeep.dataloaders import DataLoaderImbalanced\n",
    "from pytorch_widedeep.metrics import Accuracy, Precision\n",
    "from pytorch_widedeep.initializers import XavierNormal\n",
    "from pytorch_widedeep.datasets import load_bio_kdd04\n",
    "\n",
    "from sklearn.model_selection import train_test_split\n",
    "from sklearn.metrics import classification_report\n",
    "\n",
    "import time\n",
    "import datetime\n",
    "\n",
    "import warnings\n",
    "\n",
    "warnings.filterwarnings(\"ignore\", category=DeprecationWarning)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "e72ebad3",
   "metadata": {},
   "outputs": [],
   "source": [
    "df = load_bio_kdd04(as_frame=True)\n",
    "# drop columns we won't need in this example\n",
    "df.drop(columns=[\"EXAMPLE_ID\", \"BLOCK_ID\"], inplace=True)\n",
    "\n",
    "df_train, df_valid = train_test_split(\n",
    "    df, test_size=0.2, stratify=df[\"target\"], random_state=1\n",
    ")\n",
    "df_valid, df_test = train_test_split(\n",
    "    df_valid, test_size=0.5, stratify=df_valid[\"target\"], random_state=1\n",
    ")\n",
    "\n",
    "continuous_cols = df.drop(columns=[\"target\"]).columns.values.tolist()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "81cf3284",
   "metadata": {},
   "outputs": [],
   "source": [
    "# deeptabular\n",
    "tab_preprocessor = TabPreprocessor(continuous_cols=continuous_cols, scale=True)\n",
    "X_tab_train = tab_preprocessor.fit_transform(df_train)\n",
    "X_tab_valid = tab_preprocessor.transform(df_valid)\n",
    "X_tab_test = tab_preprocessor.transform(df_test)\n",
    "\n",
    "# target\n",
    "y_train = df_train[\"target\"].values\n",
    "y_valid = df_valid[\"target\"].values\n",
    "y_test = df_test[\"target\"].values"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "d202b17c",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "WideDeep(\n",
       "  (deeptabular): Sequential(\n",
       "    (0): TabMlp(\n",
       "      (cont_norm): Identity()\n",
       "      (encoder): MLP(\n",
       "        (mlp): Sequential(\n",
       "          (dense_layer_0): Sequential(\n",
       "            (0): Linear(in_features=74, out_features=64, bias=True)\n",
       "            (1): ReLU(inplace=True)\n",
       "            (2): Dropout(p=0.1, inplace=False)\n",
       "          )\n",
       "          (dense_layer_1): Sequential(\n",
       "            (0): Linear(in_features=64, out_features=32, bias=True)\n",
       "            (1): ReLU(inplace=True)\n",
       "            (2): Dropout(p=0.1, inplace=False)\n",
       "          )\n",
       "        )\n",
       "      )\n",
       "    )\n",
       "    (1): Linear(in_features=32, out_features=1, bias=True)\n",
       "  )\n",
       ")"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "deeptabular = TabMlp(\n",
    "    column_idx=tab_preprocessor.column_idx,\n",
    "    continuous_cols=tab_preprocessor.continuous_cols,\n",
    "    mlp_hidden_dims=[64, 32],\n",
    ")\n",
    "model = WideDeep(deeptabular=deeptabular)\n",
    "model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "c6262897",
   "metadata": {},
   "outputs": [],
   "source": [
    "trainer = Trainer(\n",
    "    model,\n",
    "    objective=\"binary\",\n",
    "    metrics=[Accuracy(), Precision()],\n",
    "    verbose=1,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "3c34b76f",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "epoch 1: 100%|██████████| 2074/2074 [00:10<00:00, 194.90it/s, loss=0.115, metrics={'acc': 0.959, 'prec': 0.9637}] \n",
      "valid: 100%|██████████| 258/258 [00:00<00:00, 330.04it/s, loss=0.189, metrics={'acc': 0.9147, 'prec': 0.9825}]\n",
      "predict: 100%|██████████| 456/456 [00:00<00:00, 679.61it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "              precision    recall  f1-score   support\n",
      "\n",
      "           0       1.00      0.98      0.99     14446\n",
      "           1       0.26      0.94      0.41       130\n",
      "\n",
      "    accuracy                           0.98     14576\n",
      "   macro avg       0.63      0.96      0.70     14576\n",
      "weighted avg       0.99      0.98      0.98     14576\n",
      "\n"
     ]
    }
   ],
   "source": [
    "train_dataloader = DataLoaderImbalanced(kwargs={\"oversample_mul\": 5})\n",
    "\n",
    "# in reality one should not really do this, but just to illustrate\n",
    "eval_dataloader = DataLoaderImbalanced(kwargs={\"oversample_mul\": 5})\n",
    "\n",
    "trainer.fit(\n",
    "    X_train={\"X_tab\": X_tab_train, \"target\": y_train},\n",
    "    X_val={\"X_tab\": X_tab_valid, \"target\": y_valid},\n",
    "    n_epochs=1,\n",
    "    batch_size=32,\n",
    "    train_dataloader=train_dataloader,\n",
    "    eval_dataloader=eval_dataloader,\n",
    ")\n",
    "\n",
    "preds = trainer.predict(X_tab=X_tab_test)\n",
    "print(classification_report(df_test[\"target\"].to_list(), preds))"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.15"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
